package com.cb2.algorithm.leetcode;

import java.util.HashMap;
import java.util.Map;

/**
 * <a href="https://leetcode.cn/problems/path-sum-iii/">路径总和 III(Path Sum III)</a>
 * <p>给定一个二叉树的根节点 root ，和一个整数 targetSum ，求该二叉树里节点值之和等于 targetSum 的 路径 的数目。
 * 路径 不需要从根节点开始，也不需要在叶子节点结束，但是路径方向必须是向下的（只能从父节点到子节点）。</p>
 * <p>
 * <b>示例：</b>
 * <pre>
 *     <a href="https://leetcode.cn/problems/path-sum-iii/">示例</a>
 * </pre>
 * </p>
 * <p>
 * <b>提示:</b>
 *  <ul>
 *     <li>二叉树的节点个数的范围是 [0,1000]</li>
 *     <li>-10^9 <= Node.val <= 10^9 </li>
 *     <li>-1000 <= targetSum <= 1000 </li>
 *  </ul>
 * </p>
 *
 * @author c2b
 * @see LC0112PathSum_I_S 路径总和 I(Path Sum I)
 * @see LC0113PathSum_II_M 路径总和 II(Path Sum II)
 * @see LC0437PathSum_III_M 路径总和 III(Path Sum III)
 * @since 2023/5/22 15:13
 */
public class LC0437PathSum_III_M {

    static class Solution {
        public int pathSum(TreeNode root, int targetSum) {
            return pathSumByIterator(root, targetSum);
            //return pathSumByRecursion(root, targetSum);
        }

        private int pathSumByIterator(TreeNode root, int targetSum) {
            // 一个节点的前缀和就是该节点到根之间的路径和。
            /*
                节点4的前缀和为：1 + 2 + 4 = 7
                节点8的前缀和：1 + 2 + 4 + 8 = 15
                节点9的前缀和：1 + 2 + 5 + 9 = 17
                           1
                         /  \
                        2    3
                       / \    \
                      4   5    6
                     / \   \
                    7   8   9
                我们只用遍历一次树，记录每个节点的前缀和，并查询该节点的祖先节点中符合条件的个数，将这个数量加到最终结果上。

                HashMap的key是前缀和， value是该前缀和的节点数量，记录数量是因为有出现复数路径的可能。
             */
            Map<Long, Integer> prefix = new HashMap<>();
            prefix.put(0L, 1);
            return dfs(root, prefix, 0, targetSum);
        }

        private int dfs(TreeNode currNode, Map<Long, Integer> prefix, long currPathSum, int targetSum) {
            if (currNode == null) {
                return 0;   // 空节点，满足条件路径数为0
            }
            int ret = 0;
            // 更新节点和
            currPathSum += currNode.val;

            // 从哈希表中获取能和preSum配对的前缀和个数
            ret += prefix.getOrDefault(currPathSum - targetSum, 0);
            // 将当前前缀和加入哈希表
            prefix.put(currPathSum, prefix.getOrDefault(currPathSum, 0) + 1);
            // 递归处理左右子树
            ret += dfs(currNode.left, prefix, currPathSum, targetSum);
            ret += dfs(currNode.right, prefix, currPathSum, targetSum);
            // 这个节点所在的路径都处理完了，路径方向必须是向下的（只能从父节点到子节点）
            prefix.put(currPathSum, prefix.getOrDefault(currPathSum, 0) - 1);

            return ret;
        }

        private int pathSumByRecursion(TreeNode root, int targetSum) {
            if (root == null) {
                return 0;
            }
            // 1.找出从根节点出发，所有节点值之和为 targetSum 的路径数目
            // 2.以左子节点为根节点，所有节点值之和为 targetSum 的路径数目
            // 3.以右子节点为根节点，所有节点值之和为 targetSum 的路径数目
            return dfs(root, targetSum) + pathSum(root.left, targetSum) + pathSum(root.right, targetSum);
        }

        private int dfs(TreeNode currNode, long targetSum) {
            int res = 0;
            if (currNode == null) {
                return res;
            }
            targetSum -= currNode.val;
            if (targetSum == 0) {
                ++res;
            }
            return res + dfs(currNode.left, targetSum) + dfs(currNode.right, targetSum);
        }
    }
}